#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from std_msgs.msg import Float64
from geometry_msgs.msg import PoseStamped
import tkinter as tk
from tkinter import ttk
import threading


class AltitudeController(Node):
    """高度控制GUI节点"""
    
    def __init__(self):
        super().__init__('altitude_controller')
        
        # 发布目标高度
        self.target_altitude_pub = self.create_publisher(
            Float64, 
            '/target_altitude', 
            10
        )
        
        # 订阅当前位置（用于显示当前高度）
        self.pose_sub = self.create_subscription(
            PoseStamped,
            '/drone_pose',
            self.pose_callback,
            10
        )
        
        self.current_altitude = 0.0
        self.target_altitude = 2.0  # 默认目标高度2米
        
        self.get_logger().info('高度控制节点已启动')
        
        # 在单独线程中运行GUI
        self.gui_thread = threading.Thread(target=self.create_gui)
        self.gui_thread.daemon = True
        self.gui_thread.start()
        
        # 定时发布目标高度
        self.timer = self.create_timer(0.1, self.publish_target_altitude)
    
    def pose_callback(self, msg):
        """更新当前高度"""
        self.current_altitude = msg.pose.position.z
    
    def publish_target_altitude(self):
        """发布目标高度"""
        msg = Float64()
        msg.data = self.target_altitude
        self.target_altitude_pub.publish(msg)
    
    def create_gui(self):
        """创建GUI界面"""
        self.root = tk.Tk()
        self.root.title("无人机高度控制")
        self.root.geometry("400x300")
        
        # 当前高度显示
        current_frame = ttk.Frame(self.root)
        current_frame.pack(pady=10)
        
        ttk.Label(current_frame, text="当前高度:").pack(side=tk.LEFT)
        self.current_label = ttk.Label(current_frame, text="0.00m", font=("Arial", 12, "bold"))
        self.current_label.pack(side=tk.LEFT, padx=10)
        
        # 目标高度设置
        target_frame = ttk.Frame(self.root)
        target_frame.pack(pady=20)
        
        ttk.Label(target_frame, text="目标高度 (米):").pack()
        
        # 输入框
        self.altitude_var = tk.StringVar(value="2.00")
        self.altitude_entry = ttk.Entry(
            target_frame, 
            textvariable=self.altitude_var, 
            width=10, 
            font=("Arial", 14)
        )
        self.altitude_entry.pack(pady=5)
        self.altitude_entry.bind('<Return>', self.set_altitude_from_entry)
        
        # 设置按钮
        ttk.Button(
            target_frame, 
            text="设置高度", 
            command=self.set_altitude_from_entry
        ).pack(pady=5)
        
        # 快速设置按钮
        quick_frame = ttk.Frame(self.root)
        quick_frame.pack(pady=20)
        
        ttk.Label(quick_frame, text="快速设置:").pack()
        
        button_frame = ttk.Frame(quick_frame)
        button_frame.pack(pady=5)
        
        heights = [0.5, 1.0, 1.5, 2.0, 3.0, 5.0]
        for height in heights:
            ttk.Button(
                button_frame,
                text=f"{height}m",
                command=lambda h=height: self.set_altitude(h)
            ).pack(side=tk.LEFT, padx=2)
        
        # 状态信息
        status_frame = ttk.Frame(self.root)
        status_frame.pack(pady=20)
        
        self.status_label = ttk.Label(
            status_frame, 
            text="状态: 待命", 
            font=("Arial", 10)
        )
        self.status_label.pack()
        
        # 更新显示的定时器
        self.update_display()
        
        self.root.mainloop()
    
    def set_altitude_from_entry(self, event=None):
        """从输入框设置高度"""
        try:
            altitude = float(self.altitude_var.get())
            if 0 <= altitude <= 10:  # 限制在0-10米范围
                self.set_altitude(altitude)
            else:
                self.status_label.config(text="错误: 高度必须在0-10米范围内")
        except ValueError:
            self.status_label.config(text="错误: 请输入有效的数字")
    
    def set_altitude(self, altitude):
        """设置目标高度"""
        self.target_altitude = altitude
        self.altitude_var.set(f"{altitude:.2f}")
        self.status_label.config(text=f"目标高度已设置为: {altitude:.2f}m")
        self.get_logger().info(f'目标高度设置为: {altitude:.2f}m')
    
    def update_display(self):
        """更新显示信息"""
        if hasattr(self, 'current_label'):
            self.current_label.config(text=f"{self.current_altitude:.2f}m")
        
        # 每100ms更新一次
        if hasattr(self, 'root'):
            self.root.after(100, self.update_display)


def main(args=None):
    rclpy.init(args=args)
    
    node = AltitudeController()
    
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main() 
